from typing import Iterable

import openai
from openai.types.chat import ChatCompletionMessageParam

from WAgents.llms.base import LLMMessage, LLMResponse, LLMClientConfig, LLMUsage
from WAgents.llms.base_llm import BaseLLMClient


class OpenAILLMClient(BaseLLMClient):

    @staticmethod
    def parse_messages(messages: list[LLMMessage]) -> Iterable[ChatCompletionMessageParam]:
        openai_messages = []
        for msg in messages:
            openai_messages.append({
                "role": msg.role,
                "content": msg.content,
            })
        return openai_messages

    @staticmethod
    def build_user_message(content: str) -> LLMMessage:
        return LLMMessage(role="user", content=content)

    def __init__(self, config: LLMClientConfig):
        super().__init__(config)

        self.client: openai.OpenAI = openai.OpenAI(
            api_key=config.api_key,
            base_url=config.base_url
        )


    def chat(self, messages: list[LLMMessage]) -> LLMResponse:
        completion = self.client.chat.completions.create(
            model=self.config.model,
            messages=self.parse_messages(messages),
            stream=False
        )

        return LLMResponse(
            result_type=completion.choices[0].finish_reason,
            content=completion.choices[0].message.content,
            usage=LLMUsage(
                model=completion.model,
                input_tokens=completion.usage.prompt_tokens,
                output_tokens=completion.usage.completion_tokens
            )
        )

